use std::collections::VecDeque;
use std::mem;

/// A node in the trie, containing a value and an array of child nodes.
struct TrieNode<T> {
    data: Option<T>,
    use_count: usize,
    next: [Option<Box<TrieNode<T>>>; 256],
}

impl<T> TrieNode<T> {
    /// Creates a new `TrieNode` with no data.
    fn new() -> Self {
        TrieNode {
            data: None,
            use_count: 0,
            next: Default::default(),
        }
    }
}

/// A trie data structure for mapping strings to values.
pub struct Trie<T> {
    root_node: Option<Box<TrieNode<T>>>,
}

impl<T> Trie<T> {
    /// Creates a new, empty `Trie`.
    pub fn new() -> Self {
        Trie { root_node: None }
    }

    /// Inserts a key-value pair into the trie.
    /// Returns `true` if the insertion was successful, `false` otherwise.
    pub fn insert(&mut self, key: &str, value: T) -> bool {
        if key.is_empty() {
            return false;
        }

        let mut rover = &mut self.root_node;
        for &c in key.as_bytes() {
            if rover.is_none() {
                *rover = Some(Box::new(TrieNode::new()));
            }
            let node = rover.as_mut().unwrap();
            node.use_count += 1;
            rover = &mut node.next[c as usize];
        }

        if let Some(node) = rover {
            if node.data.is_some() {
                node.data = Some(value);
                return true;
            }
        }

        if rover.is_none() {
            *rover = Some(Box::new(TrieNode::new()));
        }
        let node = rover.as_mut().unwrap();
        node.data = Some(value);
        node.use_count += 1;
        true
    }

    /// Removes a key from the trie.
    /// Returns `true` if the key was present and removed, `false` otherwise.
    pub fn remove(&mut self, key: &str) -> bool {
        if key.is_empty() {
            return false;
        }

        let mut path = VecDeque::new();
        let mut rover = &mut self.root_node;

        for &c in key.as_bytes() {
            if rover.is_none() {
                return false;
            }
            path.push_back(rover.take());
            rover = &mut rover.as_mut().unwrap().next[c as usize];
        }

        if rover.is_none() || rover.as_ref().unwrap().data.is_none() {
            return false;
        }

        rover.as_mut().unwrap().data = None;

        while let Some(mut node) = path.pop_back() {
            node.as_mut().unwrap().use_count -= 1;
            if node.as_ref().unwrap().use_count == 0 {
                *rover = None;
            } else {
                *rover = node;
                break;
            }
            rover = &mut node.as_mut().unwrap().next[key.as_bytes()[path.len()] as usize];
        }

        true
    }

    /// Looks up a key in the trie and returns its associated value.
    /// Returns `None` if the key is not found.
    pub fn lookup(&self, key: &str) -> Option<&T> {
        let mut rover = &self.root_node;
        for &c in key.as_bytes() {
            if rover.is_none() {
                return None;
            }
            rover = &rover.as_ref().unwrap().next[c as usize];
        }
        rover.as_ref().and_then(|node| node.data.as_ref())
    }

    /// Returns the number of entries in the trie.
    pub fn num_entries(&self) -> usize {
        self.root_node.as_ref().map_or(0, |node| node.use_count)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_insert_and_lookup() {
        let mut trie = Trie::new();
        assert!(trie.insert("key1", 10));
        assert!(trie.insert("key2", 20));
        assert_eq!(trie.lookup("key1"), Some(&10));
        assert_eq!(trie.lookup("key2"), Some(&20));
        assert_eq!(trie.lookup("key3"), None);
    }

    #[test]
    fn test_remove() {
        let mut trie = Trie::new();
        assert!(trie.insert("key1", 10));
        assert!(trie.insert("key2", 20));
        assert!(trie.remove("key1"));
        assert_eq!(trie.lookup("key1"), None);
        assert_eq!(trie.lookup("key2"), Some(&20));
    }

    #[test]
    fn test_num_entries() {
        let mut trie = Trie::new();
        assert_eq!(trie.num_entries(), 0);
        assert!(trie.insert("key1", 10));
        assert_eq!(trie.num_entries(), 1);
        assert!(trie.insert("key2", 20));
        assert_eq!(trie.num_entries(), 2);
        assert!(trie.remove("key1"));
        assert_eq!(trie.num_entries(), 1);
    }
}
